// FAT32Inserter.c
// main source code for the FAT32Inserter utility

/*
This utility is designed to copy boot code into the bootsector of a FAT32 formatted volume, without destroying the BPB.

*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

char *c_usage = "Usage: -v (verbose mode) -d=[Disk name] -f=[Bootsector file name]";
char *c_noargs = "Error: No arguments present.";
char *c_missmatch = "Error: Parameter miss-match.";
char *c_file_error = "Error: Could not load boot code.";
char *c_disk_error = "Error: Could not read the specified drive.";
char *print_msg = "Data written to disk:";
char *c_write_error = "Error: Could not write to disk";

typedef struct
{
	char	*disk_name;
	char	*boot_code_name;
	int	verbose;
} Arguments;

typedef struct
{
	char	*code;
	int	size;	// should be 512 on success, not 512 on error
} Sector;

Arguments *GetArguments(int argc, char **argv);
Sector *GetDiskBootSector(Arguments *arguments);
Sector *GetBootCode(Arguments *arguments);
Sector *Combine(Sector *boot_code, Sector *boot_sector);
int WriteNewBootSector(Arguments *arguments, Sector *new_boot_sector);

// the main function
int main(int argc, char **argv)
{
	// print hello message
	printf("fat32bsc tool copyright, 2012 Aidan Goddard\n");
	
	// get arguments
	Arguments *args = GetArguments(argc, argv);
	if(!args)
	{
		// error with arguments - show usage
		printf("%s\n", c_usage);
		return 1;
	}
	
	if(args->verbose)
	{
		printf("%s\n%s\n", args->disk_name, args->boot_code_name);
	}
	
	// load bootsector code into memory
	Sector *code = GetBootCode(args);
	if(!code)
	{
		// error occured loading file
		printf("%s\n", c_file_error);
		return 2;
	}
	
	// load disk's bootsector into memory
	Sector *boot = GetDiskBootSector(args);
	if(!boot)
	{
		// error occured loading bootsector
		printf("%s\n", c_disk_error);
		return 3;
	}
	
	// compine the sectors
	boot = Combine(code, boot);
	if(!boot)
	{
		return 4;
	}
	
	// write the new bootsector back to the disk
	int test = WriteNewBootSector(args, boot);
	if(test != 0)
	{
		printf("%s\n", c_write_error);
		return 5;
	}
	
	if(args->verbose)
	{
		// verbose is on, so print contents of boot->code to screen
		printf("%s\n", print_msg);
		
		for(int i = 0; i < 512; i++)
		{
			if((i % 16) == 0)
			{
				printf("\n");
			}
			
			if((unsigned char)boot->code[i] < 16)
			{
				printf(" ");
			}
			printf("%X ", (unsigned char)boot->code[i]);
		}
		
		printf("\n");
	}
	
	return 0;
}

// the function to get the arguments
Arguments *GetArguments(int argc, char **argv)
{
	// check argument count
	if(argc < 2)
	{
		printf("%s\n", c_noargs);
		return NULL;
	}
	else if(argc > 4)
	{
		printf("%s\n", c_missmatch);
		return NULL;
	}
	
	// correct number of arguments presented
	// create arguments object
	Arguments *args = (Arguments*)malloc(sizeof(Arguments));
	if(!args)
	{
		return NULL;
	}
	
	args->disk_name = (char*)malloc(128);
	if(!args->disk_name)
	{
		free(args);
		return NULL;
	}
	
	args->boot_code_name = (char*)malloc(128);
	if(!args->boot_code_name)
	{
		free(args->disk_name);
		free(args);
		return NULL;
	}
	
	
	// get arguments
	for(int i = 1; i < argc; i++)
	{
		// check for -v
		if(strncmp(argv[i], "-v", 2) == 0)
		{
			args->verbose = 1;
			continue;
		}
		
		// check for -d
		if(strncmp(argv[i], "-d=", 2) == 0)
		{
			if(strlen(argv[i]) > 3)
			{
				char *ptr = strstr(argv[i], "=");
				ptr++;
				strncpy(args->disk_name, ptr, 63);
			}
			else
			{
				printf("%s\n", c_missmatch);
			}
			continue;
		}
		
		// check for -f
		if(strncmp(argv[i], "-f=", 2) == 0)
		{
			if(strlen(argv[i]) > 3)
			{
				char *ptr = strstr(argv[i], "=");
				ptr++;
				strncpy(args->boot_code_name, ptr, 127);
			}
			else
			{
				printf("%s\n", c_missmatch);
			}
			continue;
		}
	}
	
	if(!args->disk_name)
	{
		free(args->disk_name);
		free(args->boot_code_name);
		free(args);
		return NULL;
	}
	
	if(!args->boot_code_name)
	{
		free(args->disk_name);
		free(args->boot_code_name);
		free(args);
		return NULL;
	}
	
	// return it
	return args;
}

// load boot code into memory
Sector *GetBootCode(Arguments *args)
{
	// open the file
	FILE *file = fopen(args->boot_code_name, "rb");
	
	// get the size of the file
	fseek(file, 0, SEEK_END);
	int size = ftell(file);
	
	// check the size
	if(size != 512)
	{
		return NULL;
	}
	
	// allocate a sector structure
	Sector *sector = (Sector*)malloc(sizeof(Sector));
	if(!sector)
	{
		fclose(file);
		return NULL;
	}
	
	// allocate the char array
	sector->code = (char*)malloc(512);
	if(!sector->code)
	{
		free(sector);
		fclose(file);
		return NULL;
	}
	sector->size = size;
	
	// copy the data into the array
	fseek(file, 0, SEEK_SET);
	if(fread(sector->code, 1, 512, file) != 512)
	{
		free(sector->code);
		free(sector);
		fclose(file);
		return NULL;
	}
	
	// close the file
	fclose(file);
	
	return sector;
}

// load bootsector from drive into memory
Sector *GetDiskBootSector(Arguments *args)
{
	// open disk device
	int disk = open(args->disk_name, O_RDONLY);
	if(disk < 0)
	{
		return NULL;
	}
	
	// allocate sector structure
	Sector *sector = (Sector*)malloc(sizeof(Sector));
	if(!sector)
	{
		close(disk);
		return NULL;
	}
	
	sector->code = (char*)malloc(512);
	if(!sector->code)
	{
		free(sector);
		close(disk);
		return NULL;
	}
	
	// read the code from the disk
	if(read(disk, sector->code, 512) == -1)
	{
		free(sector->code);
		free(sector);
		close(disk);
		return NULL;
	}
	
	// close the disk
	close(disk);
	
	return sector;
}

// combine the data
Sector *Combine(Sector *boot_code, Sector *boot_sector)
{
	// BPB from offset 3 - offset 90 (dec) needs copying
	
	for(int i = 3; i < 90; i++)
	{
		boot_code->code[i] = boot_sector->code[i];
	}
	
	return boot_code;
}

// write data to the disk
int WriteNewBootSector(Arguments *args, Sector *boot)
{
	// open the drive
	int disk = open(args->disk_name, O_WRONLY);
	if(disk < 0)
	{
		return 1;
	}
	
	// write the data to the drive
	if(write(disk, boot->code, 512) == -1)
	{
		close(disk);
		return 2;
	}
	
	// close the drive
	if(close(disk) != 0)
	{
		return 3;
	}
	
	
	return 0;
}


